import numpy as np

from core.function import Function
from utils.functions_conv import im2col_array


class Pooling2DWithIndexes(Function):
    def __init__(self, mpool2d):
        """
        池化操作的前向传播函数类，包括索引信息。

        Args:
            mpool2d (Pooling): 对应的前向传播池化操作函数。
        """
        self.kernel_size = mpool2d.kernel_size
        self.stride = mpool2d.stride
        self.pad = mpool2d.pad
        self.input_shpae = mpool2d.inputs[0].shape
        self.dtype = mpool2d.inputs[0].dtype
        self.indexes = mpool2d.indexes

    def forward(self, x):
        """
        前向传播，执行池化操作。

        Args:
            x (Variable): 输入变量。

        Returns:
            Variable: 池化后的输出变量。
        """
        col = im2col_array(x, self.kernel_size, self.stride, self.pad,
                           to_matrix=False)
        N, C, KH, KW, OH, OW = col.shape
        col = col.reshape(N, C, KH * KW, OH, OW)
        col = col.transpose(0, 1, 3, 4, 2).reshape(-1, KH * KW)
        indexes = self.indexes.ravel()
        col = col[np.arange(len(indexes)), indexes]
        return col.reshape(N, C, OH, OW)


